from __future__ import print_function
from numba import jit
import numpy as np
from scipy.optimize import linear_sum_assignment
from filterpy.kalman import KalmanFilter


# 计算IOU（交并比）

def calculate_iou(bb_test, bb_gt):
    """
      在两个box间计算IOU
      :param bb_test: box1 = [x1,y1,x2,y2] 左上角坐标
      :param bb_gt: box2 = [x1,y1,x2,y2] 右下角坐标
      :return: 交并比IOU
    """
    # 在两个box间的左上角坐标的最大值
    xx1 = np.maximum(bb_test[0], bb_gt[0])  # 左上角坐标x的最大值
    yy1 = np.maximum(bb_test[1], bb_gt[1])  # 左上角坐标y的最大值
    # 在两个box间的右下角坐标的最小值
    xx2 = np.minimum(bb_test[2], bb_gt[2])  # 右下角坐标x的最小值
    yy2 = np.minimum(bb_test[3], bb_test[3])  # 右下角坐标y的最小值
    # 交的宽高
    w = np.maximum(0, xx2 - xx1)
    h = np.maximum(0, yy2 - yy1)
    # 交的面积
    wh = w * h
    # 并的面积
    s = ((bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1])
         + (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1]) - wh)
    # 计算IOU并且返回IOU
    o_rate = wh / s
    return o_rate

# [x1,y1,x2,y2] 表示左上角坐标和右下角坐标，目标检测的结果以该方式表示
# [x,y,s,r]表示中心点坐标，s 是面积尺度，r是纵横比，卡尔曼滤波器中进行运动估计是使用该方式。
# 将候选框从坐标形式[x1,y1,x2,y2]转换为中心点坐标和面积的形式[x,y,s,r]
# 其中x，y是框的中心坐标，s是面积尺度，r是宽高比
def convert_bbox_to_filter_state(bbox):
    """
      将[x1,y1,x2,y2]形式的检测框转为滤波器的状态表示形式[x,y,s,r]。其中x，y是框的中心坐标，s是面积，尺度，r是宽高比
      :param bbox: [x1,y1,x2,y2] 分别是左上角坐标和右下角坐标
      :return: [ x, y, s, r ] 4行1列，其中x,y是box中心位置的坐标，s是面积，r是纵横比w/h
    """
    w = bbox[2] - bbox[0]  # 宽 x2-x1
    h = bbox[3] - bbox[1]  # 高 y2-y1
    x = bbox[0] + w / 2.0  # 检测框的中心坐标x:  x1+(x2-x1)/2.0
    y = bbox[1] + h / 2.0  # 检测框的中心坐标y: y1+(y2-y1)/2.0
    s = w * h  # 检测框的面积
    r = w / float(h)  # 检测框的宽高比
    return np.array([x, y, s, r]).reshape([4, 1])  # kalman需要四行一列的形式


# 将候选框从中心面积[x,y,s,r]的形式转换成左上角坐标和右下角坐标[x1,y1,x2,y2]的形式
# 即：将[cx，cy，s，r]的目标框表示转为[x_min，y_min，x_max，y_max]的形式
def convert_filter_state_to_bbox(x, score=None):
    """
        将[cx，cy，s，r]的目标框表示转为[x_min，y_min，x_max，y_max]的形式
        :param x:[ x, y, s, r ],其中x,y是box中心位置的坐标，s是面积，r是纵横比w/h
        :param score: 置信度
        :return:[x1,y1,x2,y2],左上角坐标和右下角坐标
        """
    w = np.sqrt(x[2] * x[3])  # w =sqrt(s*r)=sqrt(s*w/h)=sqrt(w*h * w/h)=sqrt(w*w)
    h = x[2] / w  # h =s/w =w*h/w =h
    x1 = x[0] - w / 2.0  # 左上角x坐标:x1 = x-w/2.0
    y1 = x[1] - h / 2.0  # 左上角y坐标:y1 = y-h/2.0
    x2 = x[0] + w / 2.0  # 右下角x坐标:x2 = x+w/2.0
    y2 = x[1] + h / 2.0  # 右下角y坐标:y2 = y+h/2.0
    if score is None:
        return np.array([x1, y1, x2, y2]).reshape((1, 4))
    else:
        return np.array([x1, x1, x2, y2, score]).reshape((1, 5))

if __name__ == '__main__':
    # 测试用例
    box1 = [0, 0, 4, 4]
    box2 = [0, 0, 4, 4]
    iou = calculate_iou(box1, box2)
    print(f"IoU: {iou}")

    box1 = [0, 0, 4, 4]
    box2 = [5, 5, 9, 9]
    iou = calculate_iou(box1, box2)
    print(f"IoU: {iou}")

    box1 = [0, 0, 4, 4]
    box2 = [3, 3, 7, 7]
    iou = calculate_iou(box1, box2)
    print(f"IoU: {iou}")

    # 测试用例
    box = [0, 0, 4, 4]
    filter_state = convert_bbox_to_filter_state(box)
    print(f"Filter State: {filter_state}")  # [2.0, 2.0, 16, 1.0]

    box = [3, 3, 7, 7]
    filter_state = convert_bbox_to_filter_state(box)
    print(f"Filter State: {filter_state}")  # [5.0, 5.0, 16, 1.0]

    box = [2.0, 2.0, 16, 1.0]
    bbox = convert_filter_state_to_bbox(box)
    print(f"bbox: {bbox}")  # [0, 0, 4, 4]

    box = [5.0, 5.0, 16, 1.0]
    bbox = convert_filter_state_to_bbox(box)
    print(f"bbox: {bbox}")  # [3, 3, 7, 7]
